#include "server.h"

using std::string;

#define PORTNO 12311

int ibeSERVER::HandleMessage(int clientSock, byte *szMessage)
{
  byte *szMsgString = NULL;

  if(!szMessage)
  {
    return 0;
  }
    
  switch(szMessage[0])
  {
    // Authentication (private key request)
    case 'A':
    {
      char buffer[256];
      PubPriPair *PPP = NULL;
      
      // Get the arbitrary string
      szMsgString = (byte*)malloc( (sizeof(byte)) * (strlen((const char*)szMessage)));
    
      if(!szMsgString)
      {
        printf("Could not allocate space for the ID\n");
        return -1;
      }
        
      strncpy((char*)szMsgString,(const char*)(szMessage+1),(strlen((const char*)szMessage)-1));
      szMsgString[strlen((const char*)szMessage)-1]='\0';
      
      // Get the authentication string
      memset(buffer,0,256);
		  n = read(this->clientsockfd,buffer,255);
		  
		  // Check to see if it's in the server list
      if(this->hwID.find(string((const char*)szMsgString))==hwID.end() || this->hwID[string((const char*)szMsgString)]!=string(buffer))
      {
        printf("Authentication failed for ID:%s\n",szMsgString);
        // Send the client a blank message
        byte b[4] = "";
        write(this->clientsockfd,b,4);
        free(szMsgString);
        return -1;
      }
      else
      {
        printf("ID:%s authenticated successfully!\n",szMsgString);
      }
      
      // Generate the pair
      PPP = this->cPKG->GenerateKeyPair(szMsgString);
      
      // debug
      //printf("generated key.\n");
      
      free(szMsgString);
      
      if(!PPP)
      {
        printf("Did not get public/private key information!\n");
        return -1;
      }

      // Send the private key to the client
      // Open the file
      FILE *priKey = fopen(PPP->privateFile.c_str(),"rb");
      byte *priBuffer;
      int fLen;
      
      free(PPP);
      
      if(!priKey)
      {
        //
        printf("Could not open the private key for reading!\n");
        return -1;
      }
      
      // get the file length
      fseek(priKey,0,SEEK_END);
      fLen = ftell(priKey);
      rewind(priKey);
      
      // allocate space for the file contents
      priBuffer = (byte*)malloc(fLen*sizeof(byte));
      
      if(!priBuffer)
      {
        printf("Could not allocate space for the private key buffer!\n");
        fclose(priKey);
        return -1;
      }
      
      // read the file contents (in one go)
      int n = 0;
      n = fread(priBuffer,1,fLen,priKey);

      // might as well close the file now      
      fclose(priKey);
      
      // make sure we read in everything
      if(n!=fLen)
      {
        printf("Unable to read the entire contents of the private key file!\n");
        free(priBuffer);
        return -1;
      }
      
      // send the size of the file to the client
      byte Len[25];
      sprintf((char*)Len,"%23d\n",fLen);
      //printf("len(%s)\n",Len);
      /*byte Len[4];
      Len[0] = (fLen >> 24) & 0xFF;
      Len[1] = (fLen >> 16) & 0xFF;
      Len[2] = (fLen >> 8) & 0xFF;
      Len[3] = fLen & 0xFF;*/
      //printf("length: %d%d%d%d\n",Len[0],Len[1],Len[2],Len[2]);
      
      //write(this->clientsockfd,Len,4);
      write(this->clientsockfd,Len,25);
      
      // send the file contents to the client
      write(this->clientsockfd,priBuffer,fLen);
      
      // clean up
      free(priBuffer);
      
      break;
    }
    
    // Public key request
    case 'R':
      PubPriPair *PPP = NULL;
      
      // Get the arbitrary string
      szMsgString = (byte*)malloc( (sizeof(byte)) * (strlen((const char*)szMessage)));
    
      if(!szMsgString)
      {
        printf("Could not allocate space for the ID\n");
        return -1;
      }
        
      strncpy((char*)szMsgString,(const char*)(szMessage+1),(strlen((const char*)szMessage)-1));
      szMsgString[strlen((const char*)szMessage)-1]='\0';
       
      // Generate the pair
      PPP = this->cPKG->GenerateKeyPair(szMsgString);
      
      //printf("generated key.\n");
      
      free(szMsgString);

      // Send the public key to the client
      // Open the file
      //printf("test\n");
      
     // printf("Pub key file to open: %s\n",PPP->publicFile.c_str());
      
      FILE *pubKey = fopen(PPP->publicFile.c_str(),"rb");
      byte *pubBuffer;
      int fLen;

      delete PPP;
      
      if(!pubKey)
      {
        //
        printf("Could not open the public key for reading!\n");
        return -1;
      }

      // get the file length
      fseek(pubKey,0,SEEK_END);
      fLen = ftell(pubKey);
      rewind(pubKey);

      // allocate space for the file contents
      pubBuffer = (byte*)malloc(fLen*sizeof(byte));
      
      if(!pubBuffer)
      {
        printf("Could not allocate space for the public key buffer!\n");
        fclose(pubKey);
        return -1;
      }
      // read the file contents (in one go)
      int n = 0;
      n = fread(pubBuffer,1,fLen,pubKey);
      
      // might as well close the file now      
      fclose(pubKey);
      
      // make sure we read in everything
      if(n!=fLen)
      {
        printf("Unable to read the entire contents of the public key file!\n");
        free(pubBuffer);
        return -1;
      }
      
      // send the size of the file to the client
      /*byte Len[4];
      Len[0] = (fLen >> 24) & 0xFF;
      Len[1] = (fLen >> 16) & 0xFF;
      Len[2] = (fLen >> 8) & 0xFF;
      Len[3] = fLen & 0xFF;*/
      byte Len[25];
      sprintf((char*)Len,"%23d\n",fLen);
      //printf("len(%s)\n",Len);
      
      //printf("length: %x%x%x%x(%d)\n",Len[0],Len[1],Len[2],Len[3],fLen);
      
      //write(this->clientsockfd,Len,4);
      write(this->clientsockfd,Len,25);
      
      // send the file contents to the client
      write(this->clientsockfd,pubBuffer,fLen);
      
      // clean up
      free(pubBuffer);
      
    break;
  }
    
  return 0;
}
  
ibeSERVER::ibeSERVER(byte *key)
{
  this->cPKG = new PKG(key);
}
  
int ibeSERVER::Run(void)
{
	struct sockaddr_in serv_addr, client_addr;
	byte buffer[256];
	
  // Do server stuff now
  sockfd = socket(AF_INET, SOCK_STREAM, 0);

  if(sockfd < 0)
  {
    fprintf(stderr,"Could not open socket!\n");
    return -1;
  }

	printf("Socket opened...\n");

	// set up server socket information
	memset(&serv_addr,0,sizeof(serv_addr));
	portno = PORTNO;
	serv_addr.sin_family = AF_INET;
	serv_addr.sin_addr.s_addr = INADDR_ANY;
	serv_addr.sin_port = htons(portno);
	if(bind(sockfd,(struct sockaddr*)&serv_addr,sizeof(serv_addr))<0)
	{
	  printf("Unable to bind to port!\n");
	  return -1;
	}
	printf("Bound socket to port %d...\n",portno);

	listen(sockfd,1);

	printf("Listening for a client...\n");

	clientlen = sizeof(client_addr);
	
	while(1)
	{
	  int pid=0;

	  // accept a client connection
	  clientsockfd = accept(sockfd,(struct sockaddr*)&client_addr,(socklen_t*)&clientlen);
	  if(clientsockfd < 0)
	  {
		  printf("Error on accept!\n");
		  return -1;
	  }
	  printf("Client connected...\n");

	  // split the server
	  pid = fork();

	  if(pid<0)
	  {
	  // Error
		  printf("fork failed!\n");
		  close(clientsockfd);
		  close(sockfd);
		  return -1;
	  }

	  if(pid==0)
	  {
	  // Child, only deal with the client socket now
		  close(sockfd);
		  //do key stuff
		  //for now, just read something and print it out
		  memset(buffer,0,256);
		  n = read(clientsockfd,buffer,255);
		  printf("Message from client: %s\n",buffer);
		
		  n = HandleMessage(clientsockfd,buffer);

		  close(clientsockfd);
		  return 0;
	  }
	  else
	  // Parent
		  close(clientsockfd);

	  memset(buffer,0,256);

	  //read(clientsockfd,buffer,255);
	  //write(clientsockfd,"xxx",3);
  }
}

bool ibeSERVER::LoadAuthen(const char *fiName)
{
  if(!fiName)
    return false;

  //////////////////////
	// Read id,hwid pairs
	FILE *fiAuthen = fopen(fiName,"rb");
	byte buffer[256];
	
	while(!feof(fiAuthen))
	{
	  // Read in the id/hwid pairs
	  fgets((char*)buffer,255,fiAuthen);
	  
	  if(feof(fiAuthen))
	  {
	    break;
	  }
	  
	  buffer[strlen((const char*)buffer)-1]='\0';
	  
    string str = string((const char*)buffer);
    
    int loc = str.find(' ');
    
    if(str.find(' ',loc)==string::npos)
    {
      printf("authen file invalid!\n");
	    fclose(fiAuthen);
	    return false;
    }
    
    // Separate the id from the hwid
    string sID = str.substr(0,loc);
    string shwID = str.substr(loc+1);
    
    // Add the mapping to the list
    this->hwID[sID]=shwID;
    
    // Debug
    printf("Added %s -> %s\n",sID.c_str(),shwID.c_str());
	}
	
	return true;
}

int main(int argc, char *argv[])
{	
	FILE *fiSettings = NULL;
	unsigned char *settingsContent;
	int n=0,flen=0;
	byte *key;
	
	
	///////////////
	// Read server/PKG settings
	
	// set up the PKG (we'll use a file for this)
	fiSettings = fopen("settings","rb");
	
	if(!fiSettings)
	{
	  printf("No settings file found!");
	  return -1;
	}
	
	//TODO Change this whole file reading part to use fgets
	fseek(fiSettings,0,SEEK_END);
	
	flen = ftell(fiSettings);
	
	rewind(fiSettings);
	
	settingsContent = (byte*)malloc(sizeof(byte)*flen);
	
	if(!settingsContent)
	{
	  printf("Could not allocate memory for settings contents!\n");
	  fclose(fiSettings);
	  return -1;
  }
  	
	n=fread(settingsContent,1,flen,fiSettings);
	  
	if(n<0 || n!=flen)
	{
	  printf("Error reading settings file!\n");
	  fclose(fiSettings);
	  return -1;
  }
    
	fclose(fiSettings);
	
	settingsContent[flen-1]='\0';
	
	if(strncmp("key=",(const char*)settingsContent,4)!=0)
	{
	  printf("Invalid settings file!\n");
	  //printf("file:\n\n%s\n",settingsContent);
	  free(settingsContent);
	  return -1;
	}
	
	key=(byte*)(settingsContent+4);
	printf("Using key:%s\n",key);
	
	//////////
	// Initialize the server
  ibeSERVER server(key);
  
  /////////
  // Set up authentication mapping
  if(!server.LoadAuthen("authen"))
  {
    printf("Error loading authen file!\n");
    free(settingsContent);
    return -1;
  }
  
  // Clean up memory before we start the server
  free(settingsContent);
  
  //////////////
  // Start the server
  server.Run();
	
	return 0;
}
